"""
批量图像质量检测测试脚本

递归扫描指定目录下的所有PDF、JPG、PNG文件，进行完整的图像质量检测（清晰度、分辨率、遮挡度）并生成报告
"""
import sys
from pathlib import Path
import time
import json
from typing import List, Dict, Any
from datetime import datetime
import traceback

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.services.quality.quality_checker import QualityChecker
from src.utils.converters.pdf_converter import PDFConverter
from PIL import Image
import numpy as np


class BatchQualityTester:
    """批量图像质量测试器"""
    
    def __init__(self, test_dir: str):
        self.test_dir = Path(test_dir)
        self.results = []
        self.stats = {
            'total_files': 0,
            'processed_files': 0,
            'failed_files': 0,
            'passed_count': 0,
            'warning_count': 0,
            'failed_count': 0,
            'start_time': None,
            'end_time': None
        }
        
        # 支持的文件格式
        self.supported_extensions = {'.pdf', '.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        
    def find_test_files(self) -> List[Path]:
        """递归查找所有支持的测试文件"""
        print(f"🔍 扫描目录: {self.test_dir}")
        
        if not self.test_dir.exists():
            raise FileNotFoundError(f"测试目录不存在: {self.test_dir}")
        
        test_files = []
        for ext in self.supported_extensions:
            # 递归查找所有匹配的文件
            pattern = f"**/*{ext}"
            files = list(self.test_dir.glob(pattern))
            test_files.extend(files)
            
            # 也查找大写扩展名
            pattern_upper = f"**/*{ext.upper()}"
            files_upper = list(self.test_dir.glob(pattern_upper))
            test_files.extend(files_upper)
        
        # 去重并排序
        test_files = sorted(list(set(test_files)))
        
        print(f"📁 发现 {len(test_files)} 个测试文件:")
        for i, file_path in enumerate(test_files[:10], 1):  # 只显示前10个
            rel_path = file_path.relative_to(self.test_dir)
            print(f"  {i:2d}. {rel_path}")
        
        if len(test_files) > 10:
            print(f"  ... 还有 {len(test_files) - 10} 个文件")
        
        self.stats['total_files'] = len(test_files)
        return test_files
    
    def process_single_file(self, file_path: Path) -> Dict[str, Any]:
        """处理单个文件"""
        rel_path = file_path.relative_to(self.test_dir)
        trace_id = f"batch_{int(time.time())}"
        
        result = {
            'file_path': str(rel_path),
            'full_path': str(file_path),
            'file_size': file_path.stat().st_size,
            'extension': file_path.suffix.lower(),
            'status': 'unknown',
            'error': None,
            'quality_result': None,
            'processing_time': 0
        }
        
        start_time = time.time()
        
        try:
            print(f"\n📄 处理文件: {rel_path}")
            print(f"   大小: {result['file_size']:,} bytes")
            
            # 根据文件类型处理
            if result['extension'] == '.pdf':
                # PDF文件：转换为图像
                print("   类型: PDF → 转换为图像")
                converter = PDFConverter()
                images = converter.convert_to_images(str(file_path))
                
                if not images:
                    raise ValueError("PDF转换失败，未生成图像")
                
                # 只测试第一页
                image = images[0]
                print(f"   转换成功，图像尺寸: {image.size}")
                
            else:
                # 图像文件：直接加载
                print("   类型: 图像文件")
                image = Image.open(file_path)
                print(f"   图像尺寸: {image.size}")
            
            # 转换为RGB（如果需要）
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # 完整图像质量检测（清晰度、分辨率、遮挡度）
            print("   🔍 执行图像质量检测...")
            quality_result = QualityChecker.check_image_quality(np.array(image), trace_id)
            
            result['quality_result'] = quality_result
            result['status'] = quality_result['overall_status']
            
            # 统计
            if quality_result['overall_status'] == 'passed':
                self.stats['passed_count'] += 1
                status_emoji = '✅'
            elif quality_result['overall_status'] == 'warning':
                self.stats['warning_count'] += 1
                status_emoji = '⚠️'
            else:
                self.stats['failed_count'] += 1
                status_emoji = '❌'
            
            # 显示详细结果
            clarity = quality_result['clarity']
            resolution = quality_result['resolution']
            occlusion = quality_result['occlusion']
            
            print(f"   {status_emoji} 综合评估: {quality_result['overall_status']}")
            print(f"     清晰度: {clarity['status']} (分数: {clarity['score']:.2f})")
            print(f"     分辨率: {resolution['status']} ({resolution.get('resolution', 'N/A')})")
            print(f"     遮挡度: {occlusion['status']} (分数: {occlusion.get('occlusion_score', 'N/A')})")
            
        except Exception as e:
            error_msg = str(e)
            result['error'] = error_msg
            result['status'] = 'error'
            self.stats['failed_files'] += 1
            
            print(f"   ❌ 处理失败: {error_msg}")
            
        finally:
            result['processing_time'] = time.time() - start_time
            self.stats['processed_files'] += 1
            
        return result
    
    def run_batch_test(self) -> Dict[str, Any]:
        """运行批量测试"""
        print("="*80)
        print("🚀 批量图像质量检测测试")
        print("="*80)
        
        self.stats['start_time'] = datetime.now()
        
        try:
            # 查找测试文件
            test_files = self.find_test_files()
            
            if not test_files:
                print("⚠️  未找到任何测试文件")
                return self.generate_report()
            
            print(f"\n🎯 开始处理 {len(test_files)} 个文件...")
            
            # 逐个处理文件
            for i, file_path in enumerate(test_files, 1):
                print(f"\n[{i}/{len(test_files)}]", end=" ")
                result = self.process_single_file(file_path)
                self.results.append(result)
                
                # 显示进度
                progress = (i / len(test_files)) * 100
                print(f"   进度: {progress:.1f}%")
            
        except KeyboardInterrupt:
            print("\n\n⚠️  用户中断测试")
        except Exception as e:
            print(f"\n\n❌ 测试过程中发生错误: {e}")
            traceback.print_exc()
        finally:
            self.stats['end_time'] = datetime.now()
        
        return self.generate_report()
    
    def generate_report(self) -> Dict[str, Any]:
        """生成测试报告"""
        print("\n" + "="*80)
        print("📊 测试报告")
        print("="*80)
        
        # 基本统计
        duration = 0
        if self.stats['start_time'] and self.stats['end_time']:
            duration = (self.stats['end_time'] - self.stats['start_time']).total_seconds()
        
        print(f"\n📈 统计信息:")
        print(f"  测试目录: {self.test_dir}")
        print(f"  发现文件: {self.stats['total_files']} 个")
        print(f"  处理文件: {self.stats['processed_files']} 个")
        print(f"  处理失败: {self.stats['failed_files']} 个")
        print(f"  测试耗时: {duration:.2f} 秒")
        
        # 图像质量检测结果统计
        total_tested = self.stats['passed_count'] + self.stats['warning_count'] + self.stats['failed_count']
        if total_tested > 0:
            print(f"\n🎯 图像质量检测结果:")
            print(f"  ✅ 通过 (passed):  {self.stats['passed_count']:3d} 个 ({self.stats['passed_count']/total_tested*100:.1f}%)")
            print(f"  ⚠️  警告 (warning): {self.stats['warning_count']:3d} 个 ({self.stats['warning_count']/total_tested*100:.1f}%)")
            print(f"  ❌ 失败 (failed):  {self.stats['failed_count']:3d} 个 ({self.stats['failed_count']/total_tested*100:.1f}%)")
            
            # 详细维度统计
            clarity_stats = {'passed': 0, 'warning': 0, 'failed': 0}
            resolution_stats = {'passed': 0, 'warning': 0, 'failed': 0}
            occlusion_stats = {'passed': 0, 'warning': 0, 'failed': 0}
            
            for result in self.results:
                if result['quality_result']:
                    qr = result['quality_result']
                    clarity_stats[qr['clarity']['status']] += 1
                    resolution_stats[qr['resolution']['status']] += 1
                    occlusion_stats[qr['occlusion']['status']] += 1
            
            print(f"\n📊 各维度详细统计:")
            print(f"  清晰度: ✅{clarity_stats['passed']} ⚠️{clarity_stats['warning']} ❌{clarity_stats['failed']}")
            print(f"  分辨率: ✅{resolution_stats['passed']} ⚠️{resolution_stats['warning']} ❌{resolution_stats['failed']}")
            print(f"  遮挡度: ✅{occlusion_stats['passed']} ⚠️{occlusion_stats['warning']} ❌{occlusion_stats['failed']}")
        
        # 详细结果
        if self.results:
            print(f"\n📋 详细结果:")
            print(f"{'序号':<4} {'文件路径':<40} {'综合':<8} {'清晰度':<8} {'分辨率':<8} {'遮挡度':<8} {'耗时':<8}")
            print("-" * 100)
            
            for i, result in enumerate(self.results, 1):
                file_path = result['file_path']
                if len(file_path) > 37:
                    file_path = "..." + file_path[-34:]
                
                status = result['status']
                processing_time = f"{result['processing_time']:.2f}s"
                
                if result['quality_result']:
                    qr = result['quality_result']
                    clarity_score = f"{qr['clarity']['score']:.1f}"
                    resolution_info = qr['resolution'].get('resolution', 'N/A')
                    if len(resolution_info) > 7:
                        resolution_info = resolution_info[:7]
                    occlusion_score = f"{qr['occlusion'].get('occlusion_score', 0):.3f}"
                else:
                    clarity_score = "N/A"
                    resolution_info = "N/A"
                    occlusion_score = "N/A"
                
                status_emoji = {
                    'passed': '✅',
                    'warning': '⚠️',
                    'failed': '❌',
                    'error': '💥'
                }.get(status, '❓')
                
                print(f"{i:<4} {file_path:<40} {status_emoji}{status:<7} {clarity_score:<8} {resolution_info:<8} {occlusion_score:<8} {processing_time:<8}")
        
        # 保存详细报告到文件
        # 处理datetime序列化问题
        stats_for_json = self.stats.copy()
        if stats_for_json['start_time']:
            stats_for_json['start_time'] = stats_for_json['start_time'].isoformat()
        if stats_for_json['end_time']:
            stats_for_json['end_time'] = stats_for_json['end_time'].isoformat()
        
        report_data = {
            'test_info': {
                'test_directory': str(self.test_dir),
                'test_time': stats_for_json['start_time'],
                'duration_seconds': duration
            },
            'statistics': stats_for_json,
            'results': self.results
        }
        
        # 生成报告文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        report_file = project_root / f"tests/batch_test_report_{timestamp}.json"
        
        try:
            # 自定义JSON编码器处理numpy类型
            def json_serializer(obj):
                """处理numpy和其他不可序列化的类型"""
                if hasattr(obj, 'item'):  # numpy标量
                    return obj.item()
                elif hasattr(obj, 'tolist'):  # numpy数组
                    return obj.tolist()
                elif isinstance(obj, (np.int32, np.int64, np.float32, np.float64)):
                    return obj.item()
                else:
                    return str(obj)
            
            with open(report_file, 'w', encoding='utf-8') as f:
                json.dump(report_data, f, ensure_ascii=False, indent=2, default=json_serializer)
            print(f"\n💾 详细报告已保存: {report_file}")
        except Exception as e:
            print(f"\n⚠️  保存报告失败: {e}")
        
        print("\n" + "="*80)
        print("✨ 测试完成!")
        print("="*80)
        
        return report_data


def main():
    """主函数"""
    # 默认测试目录
    default_test_dir = r"D:\公司文件\autoagensai\北交所\POC阶段\测试文件"
    
    # 解析命令行参数
    auto_run = False
    test_dir = default_test_dir
    
    for i, arg in enumerate(sys.argv[1:], 1):
        if arg in ['-y', '--yes', '--auto']:
            auto_run = True
        elif not arg.startswith('-'):
            test_dir = arg
    
    print(f"批量图像质量检测测试脚本")
    print(f"目标目录: {test_dir}")
    
    # 确认是否继续（除非自动运行）
    if not auto_run:
        try:
            response = input(f"\n是否开始测试? (y/N): ").strip().lower()
            if response not in ['y', 'yes', '是']:
                print("测试已取消")
                return
        except (KeyboardInterrupt, EOFError):
            print("\n测试已取消")
            return
    else:
        print("自动运行模式，开始测试...")
    
    # 创建测试器并运行
    tester = BatchQualityTester(test_dir)
    report = tester.run_batch_test()
    
    return report


if __name__ == "__main__":
    main()
